package fos

import scala.util.parsing.combinator.syntactical.StandardTokenParsers
import scala.util.parsing.input._
import scala.util.parsing.syntax._

/** This object implements a parser and evaluator for the
 *  simply typed lambda calculus found in Chapter 9 of
 *  the TAPL book.
 */
object SimplyTyped extends StandardTokenParsers {

  lexical.delimiters ++= List("(", ")", "\\", ".", ":", "=", "->", "{", "}", ",", "*", "+", "=>", "|")
  lexical.reserved   ++= List("Bool", "Nat", "true", "false", "if", "then", "else", "succ",
                              "pred", "iszero", "let", "in", "fst", "snd","inl","inr","as",
                              "case","of","|","=>","fix", "letrec")

  val parameterTypeMismatch = "parameter type mismatch: expected %s, found %s"
  val typeExpected = "%s type expected but %s found"
  val typeAmbiguous = "found ambiguous type, type is %s or %s"
  val variableNotInContext = "variable %s is not defined in the context"
    
    
  /** Term     ::= SimpleTerm { SimpleTerm }
   */
  def TermExp: Parser[Term] = positioned(
    rep1(SimpleTermExp) ^^ { case h::tail => (h /: tail) (Application(_,_)); case Nil/*this will never happen*/ => null } 
    | failure("illegal start of term"))

  /** SimpleTerm ::= "true"
   *               | "false"
   *               | number
   *               | "succ" Term
   *               | "pred" Term
   *               | "iszero" Term
   *               | "if" Term "then" Term "else" Term
   *               | ident
   *               | "\" ident ":" Type "." Term
   *               | "(" Term ")"
   *               | "let" ident ":" Type "=" Term "in" Term
   *               | "{" Term "," Term "}"
   *               | "fst" Term
   *               | "snd" Term
   */
  def SimpleTermExp: Parser[Term] = positioned(
       "true"          ^^^ True
     | "false"         ^^^ False
     | numericLit ^^ { case x => nvHelper(x.toInt) }
     | "succ" ~> TermExp ^^ { case t => Succ(t) }
     | "pred" ~> TermExp ^^ { case t => Pred(t) }
     | "iszero" ~> TermExp ^^ { case t => IsZero(t) }
     | ("if"~>TermExp)~("then"~>TermExp)~("else"~>TermExp)  ^^ { case t1~t2~t3 => If(t1,t2,t3) }
     | ident ^^ { case v => Variable(v) }
     | ("\\" ~> ident) ~ (":" ~> TypeExp) ~ ("." ~> TermExp) ^^ { case v ~ xt ~ t => Abstraction(Variable(v),xt,t) }
     | "(" ~> TermExp <~ ")" ^^ { case t => t }
     | ("let"~>ident)~(":"~>TypeExp)~("="~>TermExp)~("in"~>TermExp) ^^ { case id~ty~t1~t2 => Let(Variable(id),ty,t1,t2) }
     | ("{"~>TermExp<~",")~(TermExp<~"}") ^^ { case t1~t2 => Pair(t1,t2) }
     | "fst"~>TermExp ^^ { case t => First(t) }
     | "snd"~>TermExp ^^ { case t => Second(t) }
     | ("inl"~>TermExp)~("as"~>TypeExp) ^^ { case t~ty => InjectLeft(t,ty) }
     | ("inr"~>TermExp)~("as"~>TypeExp) ^^ { case t~ty => InjectRight(t,ty) }
     | ("case"~>TermExp<~"of")~("inl"~>ident)~("=>"~>TermExp)~("|"~"inr"~>ident)~("=>"~>TermExp) ^^ { case t~xl~tl~xr~tr => Case(t,Variable(xl),tl,Variable(xr),tr) }
     | "fix"~>TermExp ^^ { case t => Fix(t) }
     | ("letrec"~>ident)~(":"~>TypeExp)~("="~>TermExp)~("in"~>TermExp) ^^ { case id~ty~t1~t2 => LetRec(Variable(id),ty,t1,t2) }
     | failure("illegal start of simple term"))

  // Numerical literal decomposition
  private def nvHelper(n: Int): Term = if (n==0) Zero else Succ(nvHelper(n-1))

    
  /** Type       ::= SimpleType [ "->" Type ]
   */
  def TypeExp: Parser[Type] = positioned(
      TypeCrossSumExp ~ opt("->"~>TypeExp) ^^ { case st ~ Some(optT) => TypeFunction(st,optT); case st ~ None => st }
    | failure("illegal start of type"))

  def TypeCrossSumExp: Parser[Type] = positioned(
      SimpleTypeExp ~ opt(("*"|"+")~TypeCrossSumExp) ^^ { 
        case st ~ Some("*"~optT) => TypePair(st,optT);
        case st ~ Some("+"~optT) => TypeSum(st,optT);
        case st ~ None => st }
    | failure("illegal start of type"))
    
    
  def SimpleTypeExp: Parser[Type] = positioned(
      "Bool" ^^^ TypeBool
    | "Nat" ^^^ TypeNumeric 
    |  "(" ~> TypeExp <~ ")" ^^ { case t => t }
    | failure("illegal start of type"))

  /** Thrown when no reduction rule applies to the given term. */
  case class NoRuleApplies(t: Term) extends Exception(t.toString)

  /** Print an error message, together with the position where it occured. */
  case class TypeError(pos: Position, msg: String) extends Exception(msg) {
    override def toString = 
      msg //+ "\n" + pos.longString
  }

  /** The context is a map of variable names paired with their type. */
  type Context = Map[String, Type]

  /** Is the given term a numeric value? */
  def isNumericVal(t: Term): Boolean = t match {
    case Zero => true
    case Succ(t) => isNumericVal(t)
    case _ => false
  }
  
  /** Is the given term a value? */
  def isValue(t: Term): Boolean = t match {
    case True | False => true
    case _: Abstraction => true
    case InjectLeft(t,_) => isValue(t)
    case InjectRight(t,_) => isValue(t)
    case t => isNumericVal(t)
  }

  /** Call by value reducer. */
  def reduce(term: Term): Term = term match {
    case If(True,t2,_) => t2
    case If(False,_,t3) => t3
    case If(t1,t2,t3) => If(reduce(t1),t2,t3)
    case IsZero(Zero) => True
    case IsZero(Succ(nv)) if isNumericVal(nv) => False
    case IsZero(t) => IsZero(reduce(t))
    case Succ(t) => Succ(reduce(t))
    case Pred(Zero) => Zero
    case Pred(Succ(nv)) if isNumericVal(nv) => nv
    case Pred(t) => Pred(reduce(t))
  	case Application(t1,t2) if !isValue(t1) =>  Application(reduce(t1),t2)
  	case Application(t1,t2) if !isValue(t2) => Application(t1,reduce(t2))
    case Application(Abstraction(x, ty, t1), t2) if isValue(t2) => subst(t1, x, t2)
  	case First(Pair(v1,v2)) if isValue(v1) && isValue(v2) => v1
  	case First(t) => First(reduce(t))
  	case Second(Pair(v1,v2)) if isValue(v1) && isValue(v2) => v2
  	case Second(t) => Second(reduce(t))
  	case Pair(v,t2) if isValue(v) => Pair(v,reduce(t2)) 
  	case Pair(t1,t2) => Pair(reduce(t1),t2)
  	case Case(InjectLeft(t,_),xl,tl,_,_) if isValue(t) => subst(tl, xl, t)
  	case Case(InjectRight(t,_),_,_,xr,tr) if isValue(t) => subst(tr, xr, t)
  	case Case(t,xl,tl,xr,tr) if !isValue(t) => Case(reduce(t),xl,tl,xr,tr)
  	case InjectLeft(t,ty) if !isValue(t) => InjectLeft(reduce(t),ty)
  	case InjectRight(t,ty) if !isValue(t) => InjectRight(reduce(t),ty)
  	case Fix(Abstraction(x,ty,t)) => subst(t, x, Fix(Abstraction(x,ty,t)))
  	case Fix(t) => Fix(reduce(t)) 
    case _ => throw NoRuleApplies(term)
  }

  /** Returns the type of the given term <code>t</code>.
   *
   *  @param ctx the initial context
   *  @param t   the given term
   *  @return    the computed type
   */
  def typeof(ctx: Context, term: Term): Type = term match {
    case True | False => TypeBool
    case Zero => TypeNumeric
    case Succ(t) => typeof(ctx,t) match {
      case TypeNumeric => TypeNumeric
      case ty => throw TypeError(term.pos, typeExpected.format(TypeNumeric, ty))
    }
    case Pred(t) => typeof(ctx,t) match {
      case TypeNumeric => TypeNumeric
      case ty => throw TypeError(term.pos, typeExpected.format(TypeNumeric, ty))
    }
	case IsZero(t) => typeof(ctx,t) match {
      case TypeNumeric => TypeBool
      case ty => throw TypeError(term.pos, typeExpected.format(TypeNumeric, ty))
    }
    case If(t1,t2,t3) => (typeof(ctx,t1), typeof(ctx,t2), typeof(ctx,t3)) match {
      case (TypeBool,ty1,ty2) if ty1==ty2 => ty1
      case (TypeBool,ty1,ty2) if ty1!=ty2 => throw TypeError(term.pos, typeAmbiguous.format(ty1, ty2))
      case (ty,_,_) => throw TypeError(term.pos, typeExpected.format(TypeBool, ty))
    }
    case Abstraction(x,ty,t) => TypeFunction(ty,typeof(ctx+(x.name->ty), t))
    case Application(t1,t2) => typeof(ctx,t1) match {
      case TypeFunction(ty1,ty2) if typeof(ctx,t2)==ty1 => ty2
      case TypeFunction(ty1,_) => throw TypeError(term.pos, parameterTypeMismatch.format(ty1,typeof(ctx,t2)))
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypeFunction, tyerr))
    }
    case Pair(t1,t2) => TypePair(typeof(ctx,t1),typeof(ctx,t2))
    case First(t) => typeof(ctx,t) match {
      case TypePair(ty1,_) => ty1
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypePair, tyerr))
    }
    case Second(t) => typeof(ctx,t) match {
      case TypePair(_,ty2) => ty2
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypePair, tyerr))
    }
    case v: Variable => {
      if (ctx.contains(v.name)) ctx(v.name) 
      else throw TypeError(term.pos, variableNotInContext.format(v))
    }
    case Case(t,xl,tl,xr,tr) => typeof(ctx,t) match {
      case TypeSum(ty1,ty2) => ( typeof(ctx+(xl.name->ty1),tl), typeof(ctx+(xr.name->ty2),tr) ) match {
        case (type1,type2) if type1==type2 => type1
        case (type1,type2) => throw TypeError(term.pos, typeAmbiguous.format(type1,type2))
      }
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypeSum,tyerr))
    } 
    case InjectLeft(t,ty) => ty match {
      case TypeSum(ty1,ty2) => typeof(ctx,t) match {
      	case `ty1` => TypeSum(ty1,ty2)
      	case ty => throw TypeError(term.pos, typeExpected.format(ty1,ty))
      }
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypeSum,tyerr))
    }
    case InjectRight(t,ty) => ty match {
      case TypeSum(ty1,ty2) => typeof(ctx,t) match {
      	case `ty2` => TypeSum(ty1,ty2)
      	case ty => throw TypeError(term.pos, typeExpected.format(ty2,ty))
      }
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypeSum,tyerr))
    }
    case Fix(t) => typeof(ctx, t) match {
      case TypeFunction(ty1,ty2) if ty1==ty2 => ty1
      case TypeFunction(ty1,ty2) if ty1!=ty2 => throw TypeError(term.pos, typeExpected.format(ty1,ty2))
      case tyerr => throw TypeError(term.pos, typeExpected.format(TypeFunction,tyerr))
    }
  }

  def alpha(x: Variable, y: Variable): Variable = if (x==y) Variable(y.name+"$") else x
  
  def alpha(t: Term, y: Variable): Term = t match { 
    case True => True
    case False => False
    case Zero => Zero
    case x: Variable => alpha(x, y)
    case Abstraction(`y`,ty,t) => Abstraction(Variable(y.name+"$"), ty, subst(t,y,Variable(y.name+"$"))) 
    case Abstraction(x,ty,t) => Abstraction(x,ty, alpha(t,y)) 
    case Application(t1,t2) => Application(alpha(t1,y), alpha(t2,y))
    case IsZero(t) => IsZero(alpha(t,y))
    case If(t1,t2,t3) => If(alpha(t1,y), alpha(t2,y), alpha(t3,y))
    case Succ(t) => Succ(alpha(t,y))
    case Pred(t) => Pred(alpha(t,y))
    case Pair(t1,t2) => Pair(alpha(t1,y),alpha(t2,y))
    case First(t) => First(alpha(t,y))
    case Second(t) => Second(alpha(t,y))
    case Case(t,xl,tl,xr,tr) => Case(alpha(t,y),alpha(xl,y),alpha(tl,y),alpha(xr,y),alpha(tr,y))
    case InjectLeft(t,ty) => InjectLeft(alpha(t,y),ty)
    case InjectRight(t,ty) => InjectRight(alpha(t,y),ty)
    case Fix(t) => Fix(alpha(t, y))
  }

  def subst(v: Variable, x: Variable, s: Term): Term = if (v==x) s else v
     
  def subst(t: Term, x: Variable, s: Term): Term = t match { 
    case True => True
    case False => False
    case Zero => Zero
    case v: Variable => subst(v, x, s)
    case Abstraction(`x`, ty, t1) => Abstraction(x, ty, t1)
    case Abstraction(y,ty,t1) => if (isInFreeVarsOf(y, s)) subst(alpha(t,y),x,s) else Abstraction(y,ty, subst(t1, x, s))
    case Application(t1,t2) => Application(subst(t1,x,s), subst(t2,x,s))
    case IsZero(t) => IsZero(subst(t,x,s))
    case If(t1,t2,t3) => If(subst(t1,x,s), subst(t2,x,s), subst(t3,x,s))
    case Succ(t) => Succ(subst(t,x,s))
    case Pred(t) => Pred(subst(t,x,s))
    case Pair(t1,t2) => Pair(subst(t1,x,s),subst(t2,x,s))
    case First(t) => First(subst(t,x,s))
    case Second(t) => Second(subst(t,x,s))
    case Case(t,xl,tl,xr,tr) => Case(subst(t,x,s),xl,subst(tl,x,s),xr,subst(tr,x,s))
    case InjectLeft(t,ty) => InjectLeft(subst(t,x,s),ty)
    case InjectRight(t,ty) => InjectRight(subst(t,x,s),ty)
    case Fix(t) => Fix(subst(t,x,s)) 
  }
  
  def isInFreeVarsOf(y:Variable, s:Term): Boolean = s match { 
    case True | False | Zero => false
    case `y`  => true
    case Variable(_) => false
    case Abstraction(`y`,_,_) => false
    case Abstraction(_,_,t) => isInFreeVarsOf(y,t)
    case Application(t1,t2) => isInFreeVarsOf(y,t1) || isInFreeVarsOf(y,t2) 
    case IsZero(t) => isInFreeVarsOf(y,t)
    case If(t1,t2,t3) => isInFreeVarsOf(y,t1) || isInFreeVarsOf(y,t2) || isInFreeVarsOf(y,t3)
    case Succ(t) => isInFreeVarsOf(y,t)
    case Pred(t) => isInFreeVarsOf(y,t)
    case Pair(t1,t2) => isInFreeVarsOf(y,t1) || isInFreeVarsOf(y,t2)
    case First(t) => isInFreeVarsOf(y,t)
    case Second(t) => isInFreeVarsOf(y,t)
    case Case(t,xl,tl,xr,tr) => isInFreeVarsOf(y,t) || isInFreeVarsOf(y,tl) || isInFreeVarsOf(y,tr)
    case InjectLeft(t,ty) => isInFreeVarsOf(y,t)
    case InjectRight(t,ty) => isInFreeVarsOf(y,t)
    case Fix(t) => isInFreeVarsOf(y,t)
}
  
  /** Returns a stream of terms, each being one step of reduction.
   *
   *  @param t      the initial term
   *  @param reduce the evaluation strategy used for reduction.
   *  @return       the stream of terms representing the big reduction.
   */
  def path(t: Term, reduce: Term => Term): Stream[Term] =
    try {
      var t1 = reduce(t)
      Stream.cons(t, path(t1, reduce))
    } catch {
      case NoRuleApplies(_) =>
        Stream.cons(t, Stream.empty)
    }

  def main(args: Array[String]): Unit = {
    println("For Input:")
    val str = """
      (\x:(Nat->Nat) + Nat . case x of inl x => x (succ 2) | inr x => x) (inl \x:Nat .x as (Nat-> Nat) + Nat)
      case (inr 0 as Nat + Nat) of inl x => succ(x) | inr y => succ(succ(y))
      case (inl 0 as Nat + Nat) of inl x => succ(x) | inr y => succ(succ(y))
      (fix( \ie:(Nat->Bool) . (\x:Nat . (if iszero(x) then true else (if iszero(pred(x)) then false else ie (pred(pred x))))))) 4
      (fix \y:Nat->Nat . \x:Nat . if iszero pred x then 0 else y pred(x)) 10
      letrec y:Nat->Nat = 10 in \x:Nat . if iszero pred x then 0 else y pred(x)
      
      """.split("\n").map(_.trim).filter(_!="")
    for(arg<-str) { //TODO Delete line
     println("For Input:\n" + arg) //TODO Delete line
    val tokens = new lexical.Scanner(arg) //TODO Delete line
    //TODO Uncomment
    //val tokens = new lexical.Scanner(StreamReader(new java.io.InputStreamReader(System.in)))
    phrase(TermExp)(tokens) match {
      case Success(tree, _) =>
        try {
          println("Result:")
          println("typed: " + typeof(Map.empty, tree))
          for (t <- path(tree, reduce))
            println(t)
        } catch {
          case tperror => println(tperror.toString)
          println(tree)
        }
      case e =>
        println(e)
    }
    println
    }
  }
}
